// Copyright 2000-2017 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
package qunar.tc.decompiler.main;

import qunar.tc.decompiler.code.CodeConstants;
import qunar.tc.decompiler.main.extern.IFernflowerPreferences;
import qunar.tc.decompiler.main.rels.ClassWrapper;
import qunar.tc.decompiler.main.rels.MethodWrapper;
import qunar.tc.decompiler.modules.decompiler.exps.*;
import qunar.tc.decompiler.modules.decompiler.sforms.DirectGraph;
import qunar.tc.decompiler.modules.decompiler.stats.BasicBlockStatement;
import qunar.tc.decompiler.modules.decompiler.stats.CatchStatement;
import qunar.tc.decompiler.modules.decompiler.stats.RootStatement;
import qunar.tc.decompiler.modules.decompiler.stats.Statement;
import qunar.tc.decompiler.struct.StructField;
import qunar.tc.decompiler.struct.StructMethod;
import qunar.tc.decompiler.struct.gen.MethodDescriptor;
import qunar.tc.decompiler.struct.gen.VarType;
import qunar.tc.decompiler.util.InterpreterUtil;
import qunar.tc.decompiler.util.VBStyleCollection;

import java.util.*;
import java.util.Map.Entry;

public class ClassReference14Processor {
    private static final ExitExprent BODY_EXPR;
    private static final ExitExprent HANDLER_EXPR;

    static {
        InvocationExprent invFor = new InvocationExprent();
        invFor.setName("forName");
        invFor.setClassname("java/lang/Class");
        invFor.setStringDescriptor("(Ljava/lang/String;)Ljava/lang/Class;");
        invFor.setDescriptor(MethodDescriptor.parseDescriptor("(Ljava/lang/String;)Ljava/lang/Class;"));
        invFor.setStatic(true);
        invFor.setLstParameters(Collections.<Exprent>singletonList(new VarExprent(0, VarType.VARTYPE_STRING, null)));
        BODY_EXPR = new ExitExprent(ExitExprent.EXIT_RETURN, invFor, VarType.VARTYPE_CLASS, null);

        InvocationExprent ctor = new InvocationExprent();
        ctor.setName(CodeConstants.INIT_NAME);
        ctor.setClassname("java/lang/NoClassDefFoundError");
        ctor.setStringDescriptor("()V");
        ctor.setFunctype(InvocationExprent.TYP_INIT);
        ctor.setDescriptor(MethodDescriptor.parseDescriptor("()V"));
        NewExprent newExpr = new NewExprent(new VarType(CodeConstants.TYPE_OBJECT, 0, "java/lang/NoClassDefFoundError"), new ArrayList<Exprent>(), null);
        newExpr.setConstructor(ctor);
        InvocationExprent invCause = new InvocationExprent();
        invCause.setName("initCause");
        invCause.setClassname("java/lang/NoClassDefFoundError");
        invCause.setStringDescriptor("(Ljava/lang/Throwable;)Ljava/lang/Throwable;");
        invCause.setDescriptor(MethodDescriptor.parseDescriptor("(Ljava/lang/Throwable;)Ljava/lang/Throwable;"));
        invCause.setInstance(newExpr);
        invCause.setLstParameters(
                Collections.<Exprent>singletonList(new VarExprent(2, new VarType(CodeConstants.TYPE_OBJECT, 0, "java/lang/ClassNotFoundException"), null)));
        HANDLER_EXPR = new ExitExprent(ExitExprent.EXIT_THROW, invCause, null, null);
    }

    public static void processClassReferences(ClassesProcessor.ClassNode node) {
        // find the synthetic method Class class$(String) if present
        Map<ClassWrapper, MethodWrapper> mapClassMeths = new HashMap<>();
        mapClassMethods(node, mapClassMeths);
        if (mapClassMeths.isEmpty()) {
            return;
        }

        Set<ClassWrapper> setFound = new HashSet<>();
        processClassRec(node, mapClassMeths, setFound);

        if (!setFound.isEmpty()) {
            for (ClassWrapper wrp : setFound) {
                StructMethod mt = mapClassMeths.get(wrp).methodStruct;
                wrp.getHiddenMembers().add(InterpreterUtil.makeUniqueKey(mt.getName(), mt.getDescriptor()));
            }
        }
    }

    private static void processClassRec(ClassesProcessor.ClassNode node, final Map<ClassWrapper, MethodWrapper> mapClassMeths, final Set<? super ClassWrapper> setFound) {
        ClassWrapper wrapper = node.getWrapper();

        // search code
        for (MethodWrapper meth : wrapper.getMethods()) {
            RootStatement root = meth.root;
            if (root != null) {
                DirectGraph graph = meth.getOrBuildGraph();
                graph.iterateExprents(new DirectGraph.ExprentIterator() {
                    @Override
                    public int processExprent(Exprent exprent) {
                        for (Entry<ClassWrapper, MethodWrapper> ent : mapClassMeths.entrySet()) {
                            if (replaceInvocations(exprent, ent.getKey(), ent.getValue())) {
                                setFound.add(ent.getKey());
                            }
                        }
                        return 0;
                    }
                });
            }
        }

        // search initializers
        for (int j = 0; j < 2; j++) {
            VBStyleCollection<Exprent, String> initializers =
                    j == 0 ? wrapper.getStaticFieldInitializers() : wrapper.getDynamicFieldInitializers();

            for (int i = 0; i < initializers.size(); i++) {
                for (Entry<ClassWrapper, MethodWrapper> ent : mapClassMeths.entrySet()) {
                    Exprent exprent = initializers.get(i);
                    if (replaceInvocations(exprent, ent.getKey(), ent.getValue())) {
                        setFound.add(ent.getKey());
                    }

                    String cl = isClass14Invocation(exprent, ent.getKey(), ent.getValue());
                    if (cl != null) {
                        initializers.set(i, new ConstExprent(VarType.VARTYPE_CLASS, cl.replace('.', '/'), exprent.bytecode));
                        setFound.add(ent.getKey());
                    }
                }
            }
        }

        // iterate nested classes
        for (ClassesProcessor.ClassNode nd : node.nested) {
            processClassRec(nd, mapClassMeths, setFound);
        }
    }

    private static void mapClassMethods(ClassesProcessor.ClassNode node, Map<ClassWrapper, MethodWrapper> map) {
        boolean noSynthFlag = DecompilerContext.getOption(IFernflowerPreferences.SYNTHETIC_NOT_SET);

        ClassWrapper wrapper = node.getWrapper();

        for (MethodWrapper method : wrapper.getMethods()) {
            StructMethod mt = method.methodStruct;

            if ((noSynthFlag || mt.isSynthetic()) &&
                    mt.getDescriptor().equals("(Ljava/lang/String;)Ljava/lang/Class;") &&
                    mt.hasModifier(CodeConstants.ACC_STATIC)) {

                RootStatement root = method.root;
                if (root != null && root.getFirst().type == Statement.TYPE_TRYCATCH) {
                    CatchStatement cst = (CatchStatement) root.getFirst();
                    if (cst.getStats().size() == 2 && cst.getFirst().type == Statement.TYPE_BASICBLOCK &&
                            cst.getStats().get(1).type == Statement.TYPE_BASICBLOCK &&
                            cst.getVars().get(0).getVarType().equals(new VarType(CodeConstants.TYPE_OBJECT, 0, "java/lang/ClassNotFoundException"))) {

                        BasicBlockStatement body = (BasicBlockStatement) cst.getFirst();
                        BasicBlockStatement handler = (BasicBlockStatement) cst.getStats().get(1);

                        if (body.getExprents().size() == 1 && handler.getExprents().size() == 1) {
                            if (BODY_EXPR.equals(body.getExprents().get(0)) &&
                                    HANDLER_EXPR.equals(handler.getExprents().get(0))) {
                                map.put(wrapper, method);
                                break;
                            }
                        }
                    }
                }
            }
        }

        // iterate nested classes
        for (ClassesProcessor.ClassNode nd : node.nested) {
            mapClassMethods(nd, map);
        }
    }

    private static boolean replaceInvocations(Exprent exprent, ClassWrapper wrapper, MethodWrapper meth) {
        boolean res = false;

        while (true) {
            boolean found = false;

            for (Exprent expr : exprent.getAllExprents()) {
                String cl = isClass14Invocation(expr, wrapper, meth);
                if (cl != null) {
                    exprent.replaceExprent(expr, new ConstExprent(VarType.VARTYPE_CLASS, cl.replace('.', '/'), expr.bytecode));
                    found = true;
                    res = true;
                    break;
                }

                res |= replaceInvocations(expr, wrapper, meth);
            }

            if (!found) {
                break;
            }
        }

        return res;
    }

    private static String isClass14Invocation(Exprent exprent, ClassWrapper wrapper, MethodWrapper meth) {
        if (exprent.type == Exprent.EXPRENT_FUNCTION) {
            FunctionExprent fexpr = (FunctionExprent) exprent;
            if (fexpr.getFuncType() == FunctionExprent.FUNCTION_IIF) {
                if (fexpr.getLstOperands().get(0).type == Exprent.EXPRENT_FUNCTION) {
                    FunctionExprent headexpr = (FunctionExprent) fexpr.getLstOperands().get(0);
                    if (headexpr.getFuncType() == FunctionExprent.FUNCTION_EQ) {
                        if (headexpr.getLstOperands().get(0).type == Exprent.EXPRENT_FIELD &&
                                headexpr.getLstOperands().get(1).type == Exprent.EXPRENT_CONST &&
                                ((ConstExprent) headexpr.getLstOperands().get(1)).getConstType().equals(VarType.VARTYPE_NULL)) {

                            FieldExprent field = (FieldExprent) headexpr.getLstOperands().get(0);
                            ClassesProcessor.ClassNode fieldnode = DecompilerContext.getClassProcessor().getMapRootClasses().get(field.getClassname());

                            if (fieldnode != null && fieldnode.classStruct.qualifiedName.equals(wrapper.getClassStruct().qualifiedName)) { // source class
                                StructField fd =
                                        wrapper.getClassStruct().getField(field.getName(), field.getDescriptor().descriptorString);  // FIXME: can be null! why??

                                if (fd != null && fd.hasModifier(CodeConstants.ACC_STATIC) &&
                                        (fd.isSynthetic() || DecompilerContext.getOption(IFernflowerPreferences.SYNTHETIC_NOT_SET))) {

                                    if (fexpr.getLstOperands().get(1).type == Exprent.EXPRENT_ASSIGNMENT && fexpr.getLstOperands().get(2).equals(field)) {
                                        AssignmentExprent asexpr = (AssignmentExprent) fexpr.getLstOperands().get(1);

                                        if (asexpr.getLeft().equals(field) && asexpr.getRight().type == Exprent.EXPRENT_INVOCATION) {
                                            InvocationExprent invexpr = (InvocationExprent) asexpr.getRight();

                                            if (invexpr.getClassname().equals(wrapper.getClassStruct().qualifiedName) &&
                                                    invexpr.getName().equals(meth.methodStruct.getName()) &&
                                                    invexpr.getStringDescriptor().equals(meth.methodStruct.getDescriptor())) {

                                                if (invexpr.getLstParameters().get(0).type == Exprent.EXPRENT_CONST) {
                                                    wrapper.getHiddenMembers()
                                                            .add(InterpreterUtil.makeUniqueKey(fd.getName(), fd.getDescriptor()));  // hide synthetic field
                                                    return ((ConstExprent) invexpr.getLstParameters().get(0)).getValue().toString();
                                                }
                                            }
                                        }
                                    }
                                }
                            }
                        }
                    }
                }
            }
        }

        return null;
    }
}